# -*- coding: utf-8 -*-
import jsonpickle
from flask import g
from celery import Celery as _Celery
from flask import has_request_context, request


class Celery(_Celery):

    def init_app(self, app):
        # Instantiate celery and read config.

        celery_config = app.config['CELERY_CONFIG']

        super(Celery, self).__init__(app.import_name, broker=celery_config.get('broker_url', None))

        # Set result backend default.
        if 'result_backend' in celery_config:
            self._preconf['result_backend'] = celery_config['result_backend']

        self.conf.update(app.config['CELERY_DEFAULT_CONFIG'])
        self.conf.update(celery_config)
        self.reload_task(app)

    def reload_task(self, app):
        task_base = self.Task

        class ContextTask(task_base):

            #: Name of the additional parameter passed to tasks
            #: that contains information about the original Flask request context.
            CONTEXT_ARG_NAME = '_flask_request_context'

            def __call__(self, *args, **kwargs):
                """Execute task code with given arguments."""
                call = lambda: super(ContextTask, self).__call__(*args, **kwargs)

                context = kwargs.pop(self.CONTEXT_ARG_NAME, None)

                flask_g = kwargs.pop('__flask_g', None) or {}
                if context is None or has_request_context():
                    return call()

                with app.test_request_context(**context):
                    for k, v in jsonpickle.loads(flask_g).items():
                        setattr(g, k, v)

                    result = call()

                return result

            def apply_async(self, args=None, kwargs=None, **rest):
                self._include_request_context(kwargs)
                # arguments  = inspect.signature(super(ContextTask, self).apply_async).bind(*args, **kwargs).arguments.copy()
                kwargs['__flask_g'] = jsonpickle.dumps(g.__dict__)
                return super(ContextTask, self).apply_async(args, kwargs, **rest)

            def apply(self, args=None, kwargs=None, **rest):
                self._include_request_context(kwargs)

                return super(ContextTask, self).apply(args, kwargs, **rest)

            def retry(self, args=None, kwargs=None, **rest):
                self._include_request_context(kwargs)

                return super(ContextTask, self).retry(args, kwargs, **rest)

            def _include_request_context(self, kwargs):
                """Includes all the information about current Flask request context
                as an additional argument to the task.
                """
                if not has_request_context():
                    return

                # keys correspond to arguments of :meth:`Flask.test_request_context`
                context = {
                    'path': request.path,
                    'base_url': request.url_root,
                    'method': request.method,
                    'headers': dict(request.headers),
                    'data': request.form
                }
                if '?' in request.url:
                    context['query_string'] = request.url[(request.url.find('?') + 1):]

                if kwargs is None:
                    kwargs = {}

                kwargs[self.CONTEXT_ARG_NAME] = context

            def on_success(self, retval, task_id, args, kwargs):
                print(f'task done: {retval}')
                return super(ContextTask, self).on_success(retval, task_id, args, kwargs)

            def on_failure(self, exc, task_id, args, kwargs, einfo):
                print(f'task fail, reason: {exc}')
                return super(ContextTask, self).on_failure(exc, task_id, args, kwargs, einfo)


        setattr(ContextTask, 'abstract', True)
        setattr(self, 'Task', ContextTask)


celery = Celery('py.flask.celery')


